# -*- coding: utf-8 -*-
"""
TCN (Temporal Convolutional Network) 训练脚本
Google验证的高性能时间序列预测方法，相比LSTM速度快10倍+
"""

import pandas as pd
import numpy as np
import sys
import os
import matplotlib.pyplot as plt
from sklearn.preprocessing import RobustScaler
from sklearn.metrics import mean_squared_error
import joblib
import efinance as ef
from pypinyin import pinyin, Style

# TCN相关导入
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tcn import TCN, tcn_full_summary

def to_pinyin(text):
    s = ''
    for i in pinyin(text, style=Style.NORMAL):
        s += ''.join(i)
    return s

def configure_gpu_memory():
    """配置计算设备 - 由于RTX 5090兼容性问题，强制使用CPU"""
    # 强制使用CPU，避免RTX 5090兼容性问题
    os.environ['CUDA_VISIBLE_DEVICES'] = ''
    print("🔧 强制使用CPU模式 (RTX 5090兼容性)")
    
    # 优化CPU性能
    import multiprocessing
    cpu_count = multiprocessing.cpu_count()
    tf.config.threading.set_inter_op_parallelism_threads(cpu_count)
    tf.config.threading.set_intra_op_parallelism_threads(cpu_count)
    print(f"⚡ CPU优化: {cpu_count} 线程并行")
    
    return False

def create_tcn_model(input_shape, dilations=(1, 2, 4, 8, 16)):
    """
    创建优化的TCN模型
    
    Args:
        input_shape: 输入形状 (time_steps, features)
        dilations: 膨胀卷积的膨胀率序列 (优化后)
    
    Returns:
        Keras模型
    """
    inputs = keras.Input(shape=input_shape)
    
    # 优化的TCN层配置
    tcn_output = TCN(
        nb_filters=128,          # 增加卷积核数量 64→128
        kernel_size=3,           # 增大卷积核 2→3  
        nb_stacks=2,             # 增加堆叠层数 1→2
        dilations=dilations,     # 优化膨胀率序列
        padding='causal',        # 因果卷积
        use_skip_connections=True,  # 保持跳跃连接
        dropout_rate=0.05,       # 降低Dropout 0.1→0.05
        return_sequences=False,   # 返回最后一个时间步
        activation='relu',
        kernel_initializer='he_normal',
        use_batch_norm=True
    )(inputs)
    
    # 优化的全连接层
    x = layers.Dense(128, activation='relu', kernel_regularizer=keras.regularizers.l2(0.001))(tcn_output)
    x = layers.Dropout(0.1)(x)  # 降低Dropout
    x = layers.Dense(64, activation='relu', kernel_regularizer=keras.regularizers.l2(0.001))(x)
    x = layers.Dropout(0.1)(x)
    x = layers.Dense(32, activation='relu')(x)
    
    # 双输出：开盘价和收盘价
    outputs = layers.Dense(2, name='price_prediction')(x)
    
    model = keras.Model(inputs, outputs, name='TCN_Stock_Predictor_Optimized')
    
    # 优化的编译配置
    model.compile(
        optimizer=keras.optimizers.Adam(
            learning_rate=0.0005,  # 降低学习率 0.001→0.0005
            beta_1=0.9,
            beta_2=0.999,
            epsilon=1e-07
        ),
        loss='huber',  # 使用Huber损失，对异常值更鲁棒
        metrics=['mean_absolute_error', 'mean_squared_error']
    )
    
    return model

def prepare_tcn_data(df, time_steps=60):
    """
    准备TCN模型的训练数据 - 优化版：使用价格变化率建模
    
    Args:
        df: 包含技术指标的DataFrame
        time_steps: 时间步长
    
    Returns:
        X, y, current_prices: 特征数组、标签数组、当前价格(用于预测时转换)
    """
    # 特征列
    feature_columns = [
        'volume', 'sma5', 'sma20', 'ema12', 'ema26', 'macd', 'macdsignal',
        'macdhist', 'rsi', 'upperband', 'middleband', 'lowerband'
    ]
    
    # 当前价格列
    price_columns = ['open', 'close']
    
    # 数据预处理
    features = df[feature_columns].ffill().fillna(0)  # 修复fillna方法警告
    
    # 🚀 关键优化：计算价格变化率而非绝对价格
    current_prices = df[price_columns]
    next_prices = df[price_columns].shift(-1)
    
    # 🛠️ 安全的变化率计算，避免除零错误
    # 计算变化率：(next_price - current_price) / current_price
    # 使用更安全的方法处理除零情况
    price_change_rates = (next_prices - current_prices) / current_prices.replace(0, np.nan)
    
    # 删除包含无效值的行
    price_change_rates = price_change_rates.dropna()
    
    # 过滤极端异常值（变化率超过±100%的数据，这在正常股票中很少见）
    change_rate_mask = (np.abs(price_change_rates) < 1.0).all(axis=1)
    price_change_rates = price_change_rates[change_rate_mask]
    
    print(f"📊 数据过滤: 原始{len(df)}条 → 有效{len(price_change_rates)}条")
    
    # 确保数据长度一致，重新对齐所有数据
    min_len = min(len(features), len(price_change_rates), len(current_prices))
    # 取共同索引
    common_indices = price_change_rates.index[:min_len]
    
    features = features.loc[common_indices]
    targets = price_change_rates.loc[common_indices]
    current_prices = current_prices.loc[common_indices]
    
    print(f"📊 变化率统计: 开盘 {targets['open'].mean():.4f}±{targets['open'].std():.4f}, 收盘 {targets['close'].mean():.4f}±{targets['close'].std():.4f}")
    
    # 创建时间序列数据
    X, y, current_price_seq = [], [], []
    
    for i in range(time_steps, len(features)):
        X.append(features.iloc[i-time_steps:i].values)
        y.append(targets.iloc[i].values)
        current_price_seq.append(current_prices.iloc[i].values)  # 保存对应的当前价格
    
    return np.array(X, dtype=np.float32), np.array(y, dtype=np.float32), np.array(current_price_seq, dtype=np.float32)

# --- 1. 从命令行获取参数 ---
if len(sys.argv) < 3:
    print("错误：参数不足。")
    print("用法: python3 04d_train_tcn.py <股票代码> <frequency>")
    sys.exit(1)

stock_code = sys.argv[1]
frequency = sys.argv[2]

# --- 2. 获取公司名称并转换为拼音 ---
stock_name_pinyin = stock_code
try:
    quote_df = ef.stock.get_quote_history(stock_code, klt=101, end='20300101', limit=2)
    if not quote_df.empty and '股票名称' in quote_df.columns:
        stock_name_chinese = quote_df['股票名称'].iloc[0]
        stock_name_pinyin = to_pinyin(stock_name_chinese)
        print(f"📊 股票信息: {stock_name_chinese} ({stock_code})")
except Exception as e:
    print(f"警告：未能获取股票名称。错误: {e}")

# --- 3. 配置GPU ---
use_gpu = configure_gpu_memory()

# --- 4. 定义文件路径 ---
input_filename = f'output/data/stock_data_{stock_code}_{frequency}_with_indicators.csv'
model_path = f'output/models/tcn_model_{stock_code}_{frequency}.h5'
scaler_X_path = f'output/scalers/tcn_X_scaler_{stock_code}_{frequency}.pkl'
scaler_y_path = f'output/scalers/tcn_y_scaler_{stock_code}_{frequency}.pkl'
output_image_path = f'output/plots/prediction_vs_actual_{stock_code}_{frequency}_tcn.png'

# --- 5. 加载数据 ---
try:
    print(f"📈 开始加载处理好的数据: {input_filename}...")
    df = pd.read_csv(input_filename, index_col='date', parse_dates=True)
    print(f"✅ 数据加载完成，共 {len(df)} 条记录。")
except FileNotFoundError:
    print(f"❌ 错误: 找不到数据文件 {input_filename}。")
    print("   请先运行数据处理流程。")
    sys.exit(1)

# --- 6. 准备TCN训练数据 ---
print("🔧 准备TCN训练数据...")
print("🚀 使用价格变化率建模 - 专为高价股优化！")
TIME_STEPS = 60  # TCN的时间窗口

X, y, current_prices = prepare_tcn_data(df, TIME_STEPS)
print(f"📊 数据形状: X={X.shape}, y={y.shape}, current_prices={current_prices.shape}")

if len(X) < 100:
    print(f"⚠️  警告: 数据量较少 ({len(X)} 条)，可能影响训练效果")

# --- 7. 数据归一化 ---
print("🔧 数据归一化...")

# 特征归一化 - 重塑为2D进行归一化
scaler_X = RobustScaler()
# 将3D数据重塑为2D: (samples*timesteps, features)
X_reshaped = X.reshape(-1, X.shape[2])
X_scaled_reshaped = scaler_X.fit_transform(X_reshaped)
X_scaled = X_scaled_reshaped.reshape(X.shape)

# 目标变量归一化
scaler_y = RobustScaler()
y_scaled = scaler_y.fit_transform(y)

# 保存归一化器和价格信息
os.makedirs('output/scalers', exist_ok=True)
joblib.dump(scaler_X, scaler_X_path)
joblib.dump(scaler_y, scaler_y_path)

# 🚀 保存训练时的价格信息，用于预测时转换变化率回绝对价格
price_info_path = f'output/scalers/price_info_{stock_code}_{frequency}.pkl'
price_info = {
    'training_prices': current_prices,  # 训练时的价格序列
    'mean_prices': np.mean(current_prices, axis=0),  # 平均价格
    'std_prices': np.std(current_prices, axis=0),   # 价格标准差
}
joblib.dump(price_info, price_info_path)
print(f"💾 归一化器和价格信息已保存: {scaler_X_path}, {scaler_y_path}, {price_info_path}")

# --- 8. 划分训练集和测试集 ---
split_ratio = 0.8
split_index = int(len(X_scaled) * split_ratio)

X_train = X_scaled[:split_index]
y_train = y_scaled[:split_index]
X_test = X_scaled[split_index:]
y_test = y_scaled[split_index:]

# 对应的价格数据划分（用于后续评估）
current_prices_train = current_prices[:split_index]
current_prices_test = current_prices[split_index:]

print(f"📊 训练集: {X_train.shape[0]} 条, 测试集: {X_test.shape[0]} 条")

# --- 9. 创建TCN模型 ---
print("🏗️  创建TCN模型...")

# 优化的膨胀率序列
dilations = (1, 2, 4, 8, 16)  # 接受域 = 2 * sum(dilations) = 62，更适合60时间步
print(f"🔧 优化后膨胀率序列: {dilations}")
print(f"🔧 理论接受域: {2 * sum(dilations)} (时间步长: {TIME_STEPS})")

model = create_tcn_model(
    input_shape=(TIME_STEPS, X.shape[2]), 
    dilations=dilations
)

# 打印模型结构
print("📋 TCN模型结构:")
model.summary()

# 显示TCN特定信息
print("📊 TCN接受域信息:")
tcn_full_summary(model, expand_residual_blocks=False)

# --- 10. 训练模型 ---
print("🚀 开始训练TCN模型...")

# 优化的回调函数
callbacks = [
    keras.callbacks.EarlyStopping(
        monitor='val_mean_absolute_error',  # 监控MAE而不是loss
        patience=15,                        # 增加耐心值
        restore_best_weights=True,
        min_delta=0.0001                   # 最小改进阈值
    ),
    keras.callbacks.ReduceLROnPlateau(
        monitor='val_mean_absolute_error',
        factor=0.7,                        # 更温和的学习率衰减
        patience=8,
        min_lr=1e-6,                       # 提高最小学习率
        verbose=1
    ),
    keras.callbacks.ModelCheckpoint(
        model_path,
        monitor='val_mean_absolute_error',
        save_best_only=True,
        verbose=1
    )
]

# 训练配置优化
history = model.fit(
    X_train, y_train,
    batch_size=16,                         # 减小batch size 32→16
    epochs=150,                            # 增加epochs 100→150
    validation_data=(X_test, y_test),
    callbacks=callbacks,
    verbose=1,
    shuffle=True                           # 确保数据打乱
)

print("✅ TCN模型训练完成！")

# --- 11. 模型评估 ---
print("📊 模型评估...")
print("🚀 计算价格变化率预测的RMSE...")

# 预测变化率
y_pred_scaled = model.predict(X_test)
y_pred_change_rate = scaler_y.inverse_transform(y_pred_scaled)  # 预测的变化率
y_test_change_rate = scaler_y.inverse_transform(y_test)        # 实际的变化率

# 🚀 关键：将变化率转换回绝对价格进行评估
# 预测价格 = 当前价格 * (1 + 预测变化率)
y_pred_price = current_prices_test * (1 + y_pred_change_rate)
y_test_price = current_prices_test * (1 + y_test_change_rate)

# 计算绝对价格的RMSE（这是用户关心的指标）
rmse_open = np.sqrt(mean_squared_error(y_test_price[:, 0], y_pred_price[:, 0]))
rmse_close = np.sqrt(mean_squared_error(y_test_price[:, 1], y_pred_price[:, 1]))
rmse_avg = (rmse_open + rmse_close) / 2

# 计算变化率本身的RMSE（模型直接优化的目标）
rmse_change_open = np.sqrt(mean_squared_error(y_test_change_rate[:, 0], y_pred_change_rate[:, 0]))
rmse_change_close = np.sqrt(mean_squared_error(y_test_change_rate[:, 1], y_pred_change_rate[:, 1]))

print(f"📈 绝对价格RMSE - 开盘: {rmse_open:.4f}元, 收盘: {rmse_close:.4f}元, 平均: {rmse_avg:.4f}元")
print(f"📊 变化率RMSE - 开盘: {rmse_change_open:.6f}, 收盘: {rmse_change_close:.6f}")
print(f"📊 当前价格基准 - 开盘: {np.mean(current_prices_test[:, 0]):.2f}元, 收盘: {np.mean(current_prices_test[:, 1]):.2f}元")

# --- 12. 保存模型 ---
os.makedirs('output/models', exist_ok=True)
model.save(model_path)
print(f"💾 TCN模型已保存: {model_path}")

# --- 13. 可视化结果 ---
print("📊 生成预测对比图...")
print("🚀 绘制从变化率转换的绝对价格对比...")

# 使用英文标签避免字体问题
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(15, 10))

# 获取测试集对应的日期（简化处理）
test_dates = range(len(y_test_price))

# 开盘价对比 - 使用转换后的绝对价格
ax1.plot(test_dates, y_test_price[:, 0], label=f'{stock_name_pinyin} Actual Open Price', color='blue', linewidth=1)
ax1.plot(test_dates, y_pred_price[:, 0], label=f'{stock_name_pinyin} TCN Predicted Open Price (Change Rate Model)', color='red', linestyle='--', linewidth=1)
ax1.set_title(f'{stock_name_pinyin} ({stock_code}) - TCN Open Price Prediction via Change Rate (RMSE: {rmse_open:.4f})')
ax1.set_ylabel('Open Price (Yuan)')
ax1.legend()
ax1.grid(True, alpha=0.3)

# 收盘价对比 - 使用转换后的绝对价格
ax2.plot(test_dates, y_test_price[:, 1], label=f'{stock_name_pinyin} Actual Close Price', color='blue', linewidth=1)
ax2.plot(test_dates, y_pred_price[:, 1], label=f'{stock_name_pinyin} TCN Predicted Close Price (Change Rate Model)', color='red', linestyle='--', linewidth=1)
ax2.set_title(f'{stock_name_pinyin} ({stock_code}) - TCN Close Price Prediction via Change Rate (RMSE: {rmse_close:.4f})')
ax2.set_xlabel('Time Points')
ax2.set_ylabel('Close Price (Yuan)')
ax2.legend()
ax2.grid(True, alpha=0.3)

plt.tight_layout()
os.makedirs('output/plots', exist_ok=True)
plt.savefig(output_image_path, dpi=300, bbox_inches='tight')
print(f"📊 预测对比图已保存: {output_image_path}")

# --- 14. 训练历史可视化 ---
plt.figure(figsize=(12, 4))

plt.subplot(1, 2, 1)
plt.plot(history.history['loss'], label='Training Loss')
plt.plot(history.history['val_loss'], label='Validation Loss')
plt.title('TCN Model Loss Curve')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.grid(True)

plt.subplot(1, 2, 2)
plt.plot(history.history['mean_absolute_error'], label='Training MAE')
plt.plot(history.history['val_mean_absolute_error'], label='Validation MAE')
plt.title('TCN Model MAE Curve')
plt.xlabel('Epoch')
plt.ylabel('MAE')
plt.legend()
plt.grid(True)

training_plot_path = f'output/plots/tcn_training_history_{stock_code}_{frequency}.png'
plt.tight_layout()
plt.savefig(training_plot_path, dpi=300, bbox_inches='tight')
print(f"📊 训练历史图已保存: {training_plot_path}")

# 计算相对误差 (%)
current_price = y_test_price[-1, 1]  # 最新收盘价
rmse_percentage = (rmse_avg / current_price) * 100
mae_open = np.mean(np.abs(y_test_price[:, 0] - y_pred_price[:, 0]))
mae_close = np.mean(np.abs(y_test_price[:, 1] - y_pred_price[:, 1]))

print(f"""
🎉 优化TCN模型训练完成！

📊 性能总结:
- 平均RMSE: {rmse_avg:.4f} ({rmse_percentage:.2f}% 相对误差)
- 开盘价RMSE: {rmse_open:.4f} | MAE: {mae_open:.4f}
- 收盘价RMSE: {rmse_close:.4f} | MAE: {mae_close:.4f}
- 训练样本: {len(X_train)} 条
- 测试样本: {len(X_test)} 条
- 当前价格: {current_price:.2f}元

🔧 优化配置:
- 卷积核数量: 128 (原64)
- 网络堆叠: 2层 (原1层) 
- 学习率: 0.0005 (原0.001)
- 损失函数: Huber (原MSE)
- Batch Size: 16 (原32)

💾 生成文件:
- 模型: {model_path}
- 归一化器: {scaler_X_path}, {scaler_y_path}
- 预测图: {output_image_path}
- 训练图: {training_plot_path}

🎯 期望目标: RMSE < 0.12 (相对误差 < 0.5%)
🔄 当前状态: {'✅ 达标' if rmse_avg < 0.12 else '⚠️ 待优化'}
""")